from abc import abstractmethod
from typing import Any, Dict, List, Optional

import torch
import torch.nn as nn

from continuallearning.interfaces import (
    AdapterOutput,
    HookBasedAdapterInterface,
    TaskAwareHookManagerInterface,
)
from continuallearning.models.backbones.base import BaseBackbone

# 从专用模块导入标准化函数
from continuallearning.models.core.output_utils import standardize_adapter_output
from continuallearning.models.pefts.common.config import (
    ContextManagerConfig,
    HookManagerConfig,
)
from continuallearning.models.pefts.common.context import collect_context, task_context
from continuallearning.models.pefts.hook_managers.context_hook_manager import (
    ContextHookManager,
)

# 替换标准logging库为自定义日志系统
from continuallearning.utils.logging import get_logger

# 使用自定义日志系统创建记录器
logger = get_logger(__name__)


class BaseHookAdapter(HookBasedAdapterInterface):
    """
    Base class for hook-based adapters.

    This class extends BaseAdapter with hook functionality, providing common
    utilities for adapters that use hooks to modify model behavior.

    Args:
        backbone: Backbone model to adapt
        adapter_type: Type of adapter (PRE_BACKBONE or POST_BACKBONE)
        num_tasks: Number of tasks to support (used when task_specific=True)
    """

    _collect_features: bool

    def __init__(
        self,
        backbone: BaseBackbone,
        hook_manager_config: Dict[str, Any] | HookManagerConfig,
        context_hook_manager_config: Optional[ContextManagerConfig | Dict] = None,
    ):
        super().__init__()
        self._current_task: int = -1
        self.backbone = backbone

        # Initialize hook management
        self._hook_manager = None
        self._attached_model = None
        self._cached_features = None
        self._hook_manager_config = hook_manager_config
        self._context_hook_manager = ContextHookManager(
            self.backbone,
            context_hook_manager_config
            or ContextManagerConfig.from_dict(
                {
                    "discovery": {
                        "target_patterns": [
                            # 匹配所有编码器层 (encoder.layer.0, encoder.layer.1, ...)
                            r"^encoder\.layer\.\d+$",
                        ],
                    },
                    "extract_output": True,
                    "extract_input": False,
                },
            ),
        )

    @property
    def hook_manager(
        self,
    ) -> TaskAwareHookManagerInterface:
        """
        Get the hook manager for this adapter.

        If not already created, initializes it using the provided configuration.

        Returns:
            TaskAwareHookManagerInterface: The hook manager instance
        """
        if self._hook_manager is None:
            logger.debug(f"Creating shared hook manager for {self.__class__.__name__}")
            self._hook_manager = self._create_hook_manager(self.backbone)
        return self._hook_manager

    @abstractmethod
    def _create_hook_manager(self, model: nn.Module) -> TaskAwareHookManagerInterface:
        """
        Create a hook manager for a specific task.

        Subclasses must implement this method to create the appropriate
        hook manager for their adapter type.

        Args:
            model: Model to apply hooks to
            task_id: Task identifier

        Returns:
            TaskAwareHookManagerInterface: Hook manager for the adapter
        """
        pass

    def forward(
        self, x: torch.Tensor, task_ids: Optional[List[int]] = None, **kwargs
    ) -> AdapterOutput:
        """
        Process input through the adapter.

        Args:
            x: Input tensor
            task_id: Optional task identifier
            **kwargs: Additional arguments for the backbone

        Returns:
            AdapterOutput: Standardized output from the adapter
        """
        # Use context manager pattern for clean hook management
        if self._collect_features:
            with (
                self._context_hook_manager,
                collect_context(
                    task_id=None,
                    save_collected=True,
                ),
            ):
                _ = self.backbone(x, **kwargs)

        with (
            self,
            task_context(
                task_ids=task_ids,
                use_collected=True,
                clear_collected=True,
            ),
        ):
            try:
                # Call backbone with hooks in place
                outputs = self.backbone(x, **kwargs)
                # 使用专用函数 standardize_adapter_output 而不是通用的 standardize_output
                # 这确保返回类型始终为 AdapterOutput，符合方法声明
                std_outputs = standardize_adapter_output(outputs)
                return std_outputs
            except Exception as e:
                logger.error(f"Error during adapter forward pass: {e}")
                # Re-raise the exception with additional context
                raise RuntimeError(f"Adapter forward failed: {str(e)}") from e

    def attach(
        self,
    ) -> TaskAwareHookManagerInterface:
        """
        Attach this adapter to a model using hooks.

        This unified implementation handles both task-specific and shared adapters.

        Args:
            model: Model to apply hooks to

        Returns:
            TaskAwareHookManagerInterface: The hook manager for this adapter
        """
        hook_manager = self.hook_manager
        # Attach hooks if not already attached
        if not hook_manager.is_attached:
            hook_manager.attach()
        # Store attached model
        self._attached_model = self.backbone
        return hook_manager

    def detach(self) -> None:
        """
        Detach this adapter from the model.

        For task-specific adapters, detaches hooks for all tasks.
        """
        if self.hook_manager.is_attached:
            self.hook_manager.detach()
        self._attached_model = None

    def __enter__(self) -> "BaseHookAdapter":
        """
        Enter context manager, attaching hooks.

        Returns:
            Self: For use in with statements
        """
        if self._attached_model is None:
            if hasattr(self, "backbone"):
                self.attach()
            else:
                raise ValueError(
                    "Cannot enter context: No backbone or attached model available"
                )
        elif not self.hook_manager.is_attached:
            self.attach()
        return self

    def __exit__(
        self, exc_type: Optional[Any], exc_val: Optional[Any], exc_tb: Optional[Any]
    ) -> None:
        """
        Exit context manager, detaching hooks.

        Args:
            exc_type: Exception type if an exception was raised
            exc_val: Exception value if an exception was raised
            exc_tb: Exception traceback if an exception was raised
        """
        try:
            self.detach()
        except Exception as e:
            # Log exception but don't suppress the original exception
            logger.error(f"Error while detaching hooks in __exit__: {e}")

        # Return None to not suppress any exception that occurred
        return None

    def prepare_task(self, task_id: int, **kwargs) -> None:
        """
        Prepare adapter for the incremental training
        """
        self.register_task(task_id, **kwargs)
        self.set_task(task_id)
        self.set_trainable_tasks(task_id)

    # hook interaction
    def get_parameters(self) -> List[nn.Parameter]:
        """
        Get all parameters managed by this adapter.

        This method collects parameters from the hook manager and returns them.

        Returns:
            List[nn.Parameter]: List of parameters from all hooks
        """
        return self.hook_manager.get_parameters()

    def get_trainable_parameters(
        self,
    ) -> List[nn.Parameter]:
        return self.hook_manager.get_trainable_parameters()

    # TaskAwareInterface methods
    def register_task(
        self, task_id: int, config: Optional[Dict[str, Any]] = None
    ) -> None:
        return self.hook_manager.register_task(task_id, config)

    def set_task(self, task_id: int) -> None:
        """
        Configure the model for a specific task.

        Args:
            task_id: Task identifier
        """
        self._current_task = task_id
        self.hook_manager.set_task(task_id)

    def set_trainable_tasks(self, task_ids: int | List[int], **kwargs) -> None:
        """
        Configure the adapter for a specific task.

        This is a unified implementation for task management:
        - Updates the current task
        - For task-specific adapters, switches to task-specific parameters

        Args:
            task_id: Task identifier
        """
        self.hook_manager.set_trainable_tasks(task_ids)

    @property
    def current_task(self) -> int:
        """Get the current task ID."""
        return self._current_task

    @property
    def registered_tasks(self) -> List[int]:
        """Get list of registered task IDs."""
        return self.hook_manager.registered_tasks
